import trimesh
import json
import os
import numpy as np
from copy import deepcopy
import random
import shutil


class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)

home_dir = os.environ['HOME']

dataset_dir = '/home/v-wewei/data0/paper/grasp_suction_dataset/object_models'
dataset_dict = {}
with open('objects_dict.json', 'w', encoding='utf-8') as f:
    for root, dirs, files in os.walk(dataset_dir):
        for name in files:
            path = os.path.join(root, name)
            key = path.split('/')[-1].split('.')[0]
            #if key != '042_adjustable_wrench':
            #    continue
            #if os.path.splitext(path)[1] == '.stl':
            #    mesh = trimesh.load_mesh(path)
            #    bbox = mesh.bounding_box_oriented.primitive.extents
            #    if bbox.max() >0.25:
            #        scale = 0.25 / bbox.max()
            #        mesh.vertices *= scale
            #        mesh.export(path)

            if os.path.splitext(path)[1] == '.stl':
                print(path)
                #mesh = trimesh.load_mesh(path)
                #mesh.vertices -= mesh.center_mass
                #mesh.export(path)
                mesh = trimesh.load_mesh(path)
                print(mesh.center_mass)
                name = path.split('/')[-1].split('.')[0]
                dataset_dict['{}'.format(name)]= {}

                dataset_dict['{}'.format(name)]['path'] = path
                dataset_dict['{}'.format(name)]['bbox_oriented'] = mesh.bounding_box_oriented.primitive.extents
                dataset_dict['{}'.format(name)]['vertices'] = len(mesh.vertices)
                dataset_dict['{}'.format(name)]['faces'] = len(mesh.faces)
                dataset_dict['{}'.format(name)]['area'] = mesh.area
                mesh.density = 4000
                dataset_dict['{}'.format(name)]['mass_properties'] = mesh.mass_properties

    json.dump(dataset_dict, f, sort_keys=True, cls=NumpyEncoder, indent=4)

shutil.move('./objects_dict.json', dataset_dir)
